import tensorflow as tf
from tensorflow.keras.datasets import mnist as MNIST_dataset
from tensorflow.keras.layers import Flatten as FlattenLayer, Dense as DenseLayer
from tensorflow.keras.models import Sequential as SequentialModel
from tensorflow.keras.utils import to_categorical as to_one_hot
import numpy as np
from tabulate import tabulate as tabulate_data
import matplotlib.pyplot as plt

# 加载MNIST数据集
(train_images, train_labels), (test_images, test_labels) = MNIST_dataset.load_data()
train_images = train_images / 255.0
test_images = test_images / 255.0
train_labels = to_one_hot(train_labels)
test_labels = to_one_hot(test_labels)

# 构建神经网络模型
mnist_model = SequentialModel()  # 创建Sequential模型
mnist_model.add(FlattenLayer(input_shape=(28, 28)))  # 添加展平层
mnist_model.add(DenseLayer(20, activation='relu'))  # 添加全连接层
mnist_model.add(DenseLayer(10, activation='softmax'))  # 添加输出层

# 编译模型
mnist_model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

# 训练模型
mnist_model.fit(train_images, train_labels, epochs=20, validation_split=0.2)

# 评估模型
test_loss, test_accuracy = mnist_model.evaluate(test_images, test_labels)
print('Test accuracy:', test_accuracy)

# 保存和加载模型
model_save_path = './mnist_model5'  
loaded_mnist_model = tf.keras.models.load_model(model_save_path)